import ast
import logging
import pathlib
from typing import List, Optional, Dict, Any, Union, Tuple, Literal

from .models import ( # Use relative import
    ModuleNode, FunctionNode, MethodNode, ClassNode, VariableNode,
    ClassVariableInfo, InstanceVariableInfo, CodeKnowledgeNode, BaseNode
)
from .resolver import Resolver # Use relative import
from .utils import format_signature, get_docstring_from_node, get_value_repr # Use relative import
from ..utils.path_utils import file_path_to_module_path # Use relative import from parent

logger = logging.getLogger(__name__)

Scope = Dict[str, Any] # Placeholder for more detailed scope info later

class AstVisitor(ast.NodeVisitor):
    """
    Visits AST nodes to extract information and build CodeKnowledgeNode objects.
    Uses a Resolver instance to determine fully qualified names.
    """

    def __init__(self, project_root: pathlib.Path, file_path: pathlib.Path, resolver: Resolver):
        self.project_root: pathlib.Path = project_root
        self.file_path: pathlib.Path = file_path.resolve()
        self.resolver: Resolver = resolver
        self.module_full_path: Optional[str] = file_path_to_module_path(project_root, self.file_path)

        # Results for the current file
        self.nodes: List[CodeKnowledgeNode] = []
        # Stores imports resolved for the current module node
        self._current_module_imports: Set[str] = set()

        # Stack to keep track of the current context (module, class, function/method)
        # Each element could be a CodeKnowledgeNode or a simple dict
        self.context_stack: List[Union[CodeKnowledgeNode, BaseNode]] = []

        if self.module_full_path is not None:
            self.resolver.set_current_module(self.module_full_path, self.file_path)
            logger.info(f"Visiting: {self.file_path} (Module: {self.module_full_path})")
        else:
            # This might happen for files outside the project root or non-standard files
            logger.warning(f"Could not determine module path for file: {self.file_path}. Resolution may be impaired.")
            # Set a placeholder context if needed, or handle this case appropriately
            self.resolver.set_current_module("__unknown__", self.file_path) # Use a placeholder name


    def _add_node(self, node: CodeKnowledgeNode):
        """Adds a completed node to the results list."""
        if not node.full_path:
            logger.error(f"Attempting to add node of type {node.type} with MISSING full_path! Name: {getattr(node, 'name', 'N/A')}")
            # Optionally raise error or skip adding
            return 
        logger.debug(f"Adding node: {node.type} - {node.full_path} (from file: {self.file_path})")
        # Optional: Check for duplicates *before* adding to self.nodes (within the same file visit)
        existing = next((n for n in self.nodes if n.full_path == node.full_path), None)
        if existing:
            logger.warning(f"Duplicate full_path '{node.full_path}' detected within the *same file visit* ({self.file_path})! Existing type: {existing.type}, New type: {node.type}")
        self.nodes.append(node)

    def _get_current_context_node(self) -> Optional[Union[CodeKnowledgeNode, BaseNode]]:
        """Returns the node corresponding to the current scope (function, method, class, module)."""
        if self.context_stack:
            return self.context_stack[-1]
        return None # Should ideally not happen after visit_Module

    def _add_call(self, target_full_path: str):
        """Adds a resolved call target to the current context's 'calls' list."""
        context_node = self._get_current_context_node()
        if context_node and hasattr(context_node, 'calls') and isinstance(getattr(context_node, 'calls'), list):
            # Avoid duplicates if desired (depends on if multiple calls should be listed once)
            if target_full_path not in context_node.calls:
                context_node.calls.append(target_full_path)
                logger.debug(f"Added call '{target_full_path}' to context '{context_node.full_path}'")

    def _add_reference(self, target_full_path: str):
        """Adds a resolved reference to the current context's 'references' list."""
        context_node = self._get_current_context_node()
        if 'faker' in target_full_path.lower():
            logger.debug(f"Adding reference: Context='{context_node.full_path}', Target='{target_full_path}'")

        # Check if the reference is to the node itself (e.g., recursion, self-reference)
        if context_node and target_full_path == getattr(context_node, 'full_path', None):
            return # Don't add self-references to the list

        if context_node and hasattr(context_node, 'references') and isinstance(getattr(context_node, 'references'), list):
            if target_full_path not in context_node.references:
                context_node.references.append(target_full_path)
                logger.debug(f"Added reference '{target_full_path}' to context '{context_node.full_path}'")


    # --- Visitor Methods ---
    def visit_Module(self, node: ast.Module):
        """Entry point for visiting a module."""
        if self.module_full_path is None:
            logger.error(f"Skipping module visit due to missing module path for {self.file_path}")
            return # Cannot process if we don't know the module path

        module_node = ModuleNode(
            full_path=self.module_full_path,
            file_path=str(self.file_path),
            docstring=ast.get_docstring(node, clean=False)
        )
        self.context_stack.append(module_node)
        self._current_module_imports = set() # Reset imports for this module

        # Visit body - imports first might help resolver context, but order usually doesn't matter for static
        self.generic_visit(node) # Visit imports, functions, classes, assignments

        # Finalize module node after visiting body
        module_node.imports = sorted(list(self._current_module_imports))
        # Populate functions, classes, variables lists based on nodes added
        for n in self.nodes:
            if n.type == 'function' and n.module_full_path == self.module_full_path:
                module_node.functions.append(n.full_path)
            elif n.type == 'class' and n.module_full_path == self.module_full_path:
                module_node.classes.append(n.full_path)
            elif n.type == 'variable' and n.module_full_path == self.module_full_path:
                module_node.variables.append(n.full_path)

        # Pop context and add the completed module node
        self.context_stack.pop()
        self._add_node(module_node) # Add module node last

    def visit_Import(self, node: ast.Import):
        logger.debug(f"Visiting Import: {ast.dump(node)}")
        resolved_imports = self.resolver.resolve_import_statement(node)
        for _, _, source_module_path in resolved_imports:
            if source_module_path:
                self._current_module_imports.add(source_module_path)
        self.generic_visit(node) # Continue visit in case needed, though unlikely for Import

    def visit_ImportFrom(self, node: ast.ImportFrom):
        logger.debug(f"Visiting ImportFrom: {ast.dump(node)}")
        resolved_imports = self.resolver.resolve_from_import_statement(node)
        # Add the source module path from which names were imported
        # If resolved_imports is not empty, they all share the same source_module_path (third element)
        if resolved_imports:
            source_module_path = resolved_imports[0][2]
            if source_module_path and not source_module_path.endswith(".*"): # Don't add the wildcard marker itself
                self._current_module_imports.add(source_module_path)
        self.generic_visit(node)

    def visit_FunctionDef(self, node: ast.FunctionDef):
        func_name = node.name
        parent_context_node = self._get_current_context_node()

        # --- Check for Nested Function ---
        is_nested = isinstance(parent_context_node, (FunctionNode, MethodNode))
        if is_nested:
            if parent_context_node: # Should always exist if is_nested is True
                # Construct a conceptual path/identifier for the nested function
                nested_func_id = f"{parent_context_node.full_path}.<locals>.{func_name}"
                logger.debug(f"Detected nested function '{func_name}' inside '{parent_context_node.full_path}'. Adding ID '{nested_func_id}' to parent's nested_functions.")
                # Add to parent's list
                if hasattr(parent_context_node, 'nested_functions') and isinstance(parent_context_node.nested_functions, list):
                    if nested_func_id not in parent_context_node.nested_functions:
                        parent_context_node.nested_functions.append(nested_func_id)
                else:
                    logger.warning(f"Parent context node {type(parent_context_node)} for nested function {func_name} lacks 'nested_functions' list.")

            # >>> Option A: Do not create a separate node and do not visit the body <<<
            return # Stop processing this nested function further
            # >>> End Option A <<<
        # --- End Check for Nested Function ---

        # --- Original logic for top-level functions and methods ---
        # (This part only runs if not nested)
        full_path = self.resolver._get_full_path_for_definition(func_name)

        if not full_path:
            logger.warning(f"Could not determine full_path for function/method '{func_name}' in {self.file_path}, skipping.")
            return

        self.resolver.add_local_definition(func_name, full_path)

        signature, return_hint = format_signature(node)
        docstring = ast.get_docstring(node, clean=False)
        try:
            implementation = ast.unparse(node)
        except Exception:
            logger.warning(f"Could not unparse implementation for {full_path}", exc_info=True)
            implementation = None

        is_method = isinstance(parent_context_node, ClassNode)
        method_context_info = None
        node_to_process: Union[MethodNode, FunctionNode]

        if is_method:
            # --- Handle Method ---
            class_node = parent_context_node
            # Directly construct the full path for the method
            full_path = f"{class_node.full_path}.{func_name}"
            logger.debug(f"Method detected: Generating path directly: '{full_path}'")

            # *** CRITICAL: Do NOT add method to module-level resolver definitions ***
            # self.resolver.add_local_definition(func_name, full_path) # <-- REMOVE/COMMENT OUT THIS LINE

            # Determine method type
            method_type: Literal['instance', 'class', 'static'] = 'instance'
            decorator_names = set()
            # (Keep decorator check logic as before)
            for decorator_node in node.decorator_list:
                # Try to get simple decorator names like @classmethod, @staticmethod
                if isinstance(decorator_node, ast.Name):
                    decorator_names.add(decorator_node.id)
            if 'classmethod' in decorator_names: method_type = 'class'
            elif 'staticmethod' in decorator_names: method_type = 'static'

            # Create the MethodNode
            node_to_process = MethodNode(
                name=func_name,
                full_path=full_path,
                class_name=class_node.name,
                class_full_path=class_node.full_path,
                method_type=method_type,
                signature=signature,
                return_type_hint=return_hint,
                docstring=docstring,
                implementation=implementation,
                # decorators=[], # TODO: Resolve decorator full_paths and store them
            )

            # Add method full_path to parent class's list
            if full_path not in class_node.methods:
                class_node.methods.append(full_path)

            # Prepare context info for the resolver
            method_context_info = {
                'class_full_path': class_node.full_path,
                'is_static': (method_type == 'static'),
                'is_classmethod': (method_type == 'class'),
            }

        else: # Standalone function
            if not self.module_full_path:
                logger.error(f"Cannot determine full_path for function '{func_name}' as module path is unknown.")
                return # Cannot proceed without module path

            # Use resolver for function path generation
            full_path = self.resolver._get_full_path_for_definition(func_name)
            if not full_path:
                logger.warning(f"Could not determine full_path for function '{func_name}' in {self.file_path}, skipping.")
                return

            logger.debug(f"Function detected: Path from resolver: '{full_path}'")

            # *** Add function definition to module-level resolver ***
            self.resolver.add_local_definition(func_name, full_path)

            # Create FunctionNode
            node_to_process = FunctionNode(
                name=func_name,
                full_path=full_path,
                module_full_path=self.module_full_path,
                signature=signature,
                return_type_hint=return_hint,
                docstring=docstring,
                implementation=implementation,
                # decorators=[], # TODO: Resolve decorator full_paths and store them
            )

        # --- Processing Body (Common logic for Methods and Functions) ---
        # Ensure node_to_process and full_path are valid before proceeding
        if not node_to_process or not full_path:
            logger.error(f"Internal error: node_to_process or full_path not set for {func_name}. Aborting visit.")
            return

        self.context_stack.append(node_to_process)

        # 2. If it's a method, set the resolver's method context
        if is_method and method_context_info:
            self.resolver.set_method_context(method_context_info)

        try:
            # 3. Visit the function/method body recursively
            # This happens *only* for non-nested functions/methods with Option A
            self.generic_visit(node)
        finally:
            # 4. IMPORTANT: Clear the resolver's method context *after* visiting the body
            if is_method:
                self.resolver.set_method_context(None)

            # 5. Pop the function/method node from the visitor's context stack
            self.context_stack.pop()

        # 6. Add the completed node (with populated calls/references from body visit) to results
        self._add_node(node_to_process)

    # Handle async functions similarly
    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_ClassDef(self, node: ast.ClassDef):
        class_name = node.name
        full_path = self.resolver._get_full_path_for_definition(class_name)

        if not full_path:
            logger.warning(f"Could not determine full_path for class '{class_name}' in {self.file_path}, skipping.")
            return
        if not self.module_full_path:
            logger.error(f"Cannot create ClassNode for {class_name} as module path is unknown.")
            return

        # Add definition to resolver *before* visiting body
        self.resolver.add_local_definition(class_name, full_path)

        docstring = ast.get_docstring(node, clean=False)
        inherits_from: List[str] = []
        logger.debug(f"visit_ClassDef: Resolving base classes for: {full_path}")
        for base in node.bases:
            base_name_str = ast.unparse(base) # Get text representation for logging
            logger.debug(f"visit_ClassDef: Attempting to resolve base '{base_name_str}', type={type(base)}") # Log type
            # Resolve base class full_path
            base_full_path = self._resolve_node_to_fullpath(base)
            logger.debug(f"visit_ClassDef: Resolution result for '{base_name_str}': {base_full_path}")
            if base_full_path:
                inherits_from.append(base_full_path)
            else:
                # Log unresolved base, maybe store original name?
                unresolved_base_name = ast.unparse(base)
                logger.warning(f"Could not resolve base class '{unresolved_base_name}' for class '{full_path}'")


        class_node = ClassNode(
            name=class_name,
            full_path=full_path,
            module_full_path=self.module_full_path,
            docstring=docstring,
            inherits_from=inherits_from,
            # decorators=[], # TODO: Resolve decorator full_paths
        )

        self.context_stack.append(class_node)
        # Visit methods, nested classes, class variables defined in the body
        # Important: visit body *after* pushing context, *before* adding the node
        self.generic_visit(node)
        self.context_stack.pop()
        self._add_node(class_node)

    def visit_Assign(self, node: ast.Assign):
        # Handle module-level and class-level variable assignments
        current_context = self._get_current_context_node()
        value_repr = get_value_repr(node.value) # Use util function

        for target in node.targets:
            if isinstance(target, ast.Name):
                var_name = target.id
                if isinstance(current_context, ModuleNode):
                    full_path = self.resolver._get_full_path_for_definition(var_name)
                    if full_path:
                        self.resolver.add_local_definition(var_name, full_path)
                        var_node = VariableNode(
                            name=var_name,
                            full_path=full_path,
                            module_full_path=self.module_full_path,
                            value_repr=value_repr
                            # type_hint and docstring from AnnAssign or type comments later
                        )
                        self._add_node(var_node)
                elif isinstance(current_context, ClassNode):
                    class_node = current_context
                    full_path = f"{class_node.full_path}.{var_name}"
                    # Assume it's a class variable if assigned directly in class body
                    class_var_info = ClassVariableInfo(
                        name=var_name,
                        full_path=full_path,
                        value_repr=value_repr
                        # type_hint and docstring later
                    )
                    # Avoid duplicates if re-assigned? Or update? For now, just add.
                    class_node.class_variables.append(class_var_info)
                    # Also add to resolver for potential references within the class body
                    # self.resolver.add_local_definition(var_name, full_path) # Need scoped definitions
            elif isinstance(target, ast.Attribute):
                # Could be assignment to self.x (instance var) or cls.y (class var) inside a method
                # Need more context (am I inside __init__? is base 'self' or 'cls'?)
                # Defer detailed instance/class var inference from methods for later stage.
                pass

        # Continue visiting the value part for calls/references
        self.visit(node.value)

    def visit_AnnAssign(self, node: ast.AnnAssign):
        # Similar to Assign, but with type annotations
        current_context = self._get_current_context_node()
        type_hint = ast.unparse(node.annotation)
        value_repr = get_value_repr(node.value) if node.value else None

        if isinstance(node.target, ast.Name):
            var_name = node.target.id
            if isinstance(current_context, ModuleNode):
                full_path = self.resolver._get_full_path_for_definition(var_name)
                if full_path:
                    self.resolver.add_local_definition(var_name, full_path)
                    var_node = VariableNode(
                        name=var_name,
                        full_path=full_path,
                        module_full_path=self.module_full_path,
                        value_repr=value_repr,
                        type_hint=type_hint
                        # docstring from type comments later
                    )
                    self._add_node(var_node)
            elif isinstance(current_context, ClassNode):
                class_node = current_context
                full_path = f"{class_node.full_path}.{var_name}"
                class_var_info = ClassVariableInfo(
                    name=var_name,
                    full_path=full_path,
                    value_repr=value_repr,
                    type_hint=type_hint
                )
                class_node.class_variables.append(class_var_info)
                # self.resolver.add_local_definition(var_name, full_path) # Need scoped definitions

        # Continue visiting value and annotation for calls/references
        if node.value:
            self.visit(node.value)
        self.visit(node.annotation)

    def visit_Call(self, node: ast.Call):
        is_method = isinstance(self._get_current_context_node(), (MethodNode, FunctionNode)) # Basic check
        # Resolve the function being called
        func_full_path = self._resolve_node_to_fullpath(node.func, is_method_scope=is_method) # Pass flag

        if func_full_path:
            # Check if it's a call to super() - special handling?
            # if isinstance(node.func, ast.Name) and node.func.id == 'super':
            #     pass # Add special marker or resolve based on MRO? Defer for now.

            self._add_call(func_full_path)

        # Visit arguments and keywords as they might contain further calls/references
        for arg in node.args:
            self.visit(arg)
        for kwarg in node.keywords:
            self.visit(kwarg.value)
        # Also visit the function node itself (e.g., if it's an attribute access like obj.method())
        self.visit(node.func)


    def visit_Name(self, node: ast.Name):
        # If loading a name (reading its value), resolve it and add as reference
        if isinstance(node.ctx, ast.Load):
            is_method = isinstance(self._get_current_context_node(), (MethodNode, FunctionNode))
            resolved_path = self.resolver.resolve(node.id, is_method_scope=is_method) # Pass flag
            if resolved_path:
                # Avoid adding references to the immediate parent function/method/class itself
                current_context = self._get_current_context_node()
                if not (current_context and resolved_path == getattr(current_context, 'full_path', None)):
                    self._add_reference(resolved_path)
        # We don't need to visit children of Name
        # self.generic_visit(node) # Avoid this for Name

    def visit_Attribute(self, node: ast.Attribute):
        # If loading an attribute (reading value), resolve and add reference
        if isinstance(node.ctx, ast.Load):
            is_method = isinstance(self._get_current_context_node(), (MethodNode, FunctionNode))
            full_path = self._resolve_node_to_fullpath(node, is_method_scope=is_method) # Pass flag
            if full_path:
                # Avoid adding references to the immediate parent function/method/class itself
                current_context = self._get_current_context_node()
                if not (current_context and full_path == getattr(current_context, 'full_path', None)):
                    self._add_reference(full_path)

        # Important: Visit the base of the attribute (node.value)
        # e.g., for 'a.b.c', we need to visit 'a.b' first
        self.visit(node.value)
        # We don't visit node.attr (it's just a string)

    def _resolve_node_to_fullpath(self, node: ast.expr, *, is_method_scope: bool = False) -> Optional[str]:
        """Helper to resolve various AST expression nodes to a full_path."""
        logger.debug(f"_resolve_node: Called for node type {type(node)}, code='{ast.unparse(node)}'") # Log entry
        if isinstance(node, ast.Name):
            logger.debug(f"_resolve_node: Handling Name '{node.id}', calling resolver.resolve...")
            result = self.resolver.resolve(node.id, is_method_scope=is_method_scope)
            logger.debug(f"_resolve_node: resolver.resolve('{node.id}') returned: {result}")
            return result
        elif isinstance(node, ast.Attribute):
            logger.debug(f"_resolve_node: Handling Attribute '{node.attr}', calling resolver.resolve_attribute...")
            result = self.resolver.resolve_attribute(node.value, node.attr, is_method_scope=is_method_scope)
            logger.debug(f"_resolve_node: resolver.resolve_attribute returned: {result}")
            return result
        elif isinstance(node, ast.Subscript):
            # list[int], dict[str, Any] etc. - resolve the base value
            return self._resolve_node_to_fullpath(node.value)
        elif isinstance(node, ast.Call):
            # Calling a function that returns something - we don't know the return type easily
            # We could try resolving the function being called itself?
            # For now, return None, as we can't easily determine the path of the result.
            logger.debug(f"Cannot resolve full_path from Call node: {ast.unparse(node)}")
            return None
        # Add other node types as needed (e.g., BinOp, Constant) - unlikely to resolve to paths
        else:
            logger.debug(f"Cannot resolve node type {type(node)} to full_path: {ast.unparse(node)}")
            return None